import os, sys
import numpy as np
import scipy as sp
import matplotlib.pylab as plt

from dcds_mask import DCDS_Mask


def getWeightsOneFrame(OS_wave_dir, 
                       nsample_per_pixel=100, 
                       file_num=90, 
                       n_ref=37, 
                       n_sig=37, 
                       mask=None, 
                       out_fname=None ):
    """
    Optimize the DCDS weights from one frame.
    """

    if mask is None:
        print('mask is None, quit.')
        sys.exit(0)
    
    if np.sum(mask) != 0:
        print('sum(mask) dose not equal to zero, quit now!')
        sys.exit(0)

    pixel_num = None

    # 1) load waveform data:
    d = None
    dcnt=0
    for i in range(file_num):
        wfile = os.path.join(OS_wave_dir,str(i+1)+'_1.txt')
        if os.path.exists(wfile):
            waveform = np.loadtxt(wfile)
            pixel_num = len(waveform) // nsample_per_pixel
            waveform = waveform.reshape(pixel_num,nsample_per_pixel)
            
            if dcnt==0:
                d = waveform
            else:
                d = np.vstack((d,waveform))
        
            dcnt += 1
        else:
            print('### file : {} does not exist.'.format(wfile))
            continue
    
    idx = mask != 0
    d = d[:,idx]
    
    # 2) prepare matrix N
    N = np.identity(file_num*pixel_num)
    N = N - 1./(file_num*pixel_num)
    
    # 3) prepare matrix A
    A = np.matmul(N,d)
    AtA = np.matmul(A.transpose(), A)
    eA = np.zeros((AtA.shape[0]+2, AtA.shape[1]+2))
    eA[0:AtA.shape[0],0:AtA.shape[1]] = AtA

    i = n_ref + n_sig + 0
    for j in range(n_ref):
        eA[i,j] = 1
        eA[j,i] = 1
    
    i = n_ref + n_sig + 1
    for j in range(n_ref,n_ref+n_sig):
        eA[i,j] = 1
        eA[j,i] = 1

    # 4) prepare vector b
    b = np.zeros(n_ref+n_sig+2)
    b[n_ref+n_sig+0] = 1
    b[n_ref+n_sig+1] = -1
    
    w = np.linalg.solve(eA, b)

    if out_fname is not None:
        wout = []
        w_cnt = 0
        for i in range(nsample_per_pixel):
            if mask[i] != 0:
                wout.append(w[w_cnt])
                w_cnt += 1
            else:
                wout.append(0)
        
        wout = np.array(wout)
        np.savetxt(out_fname, wout, fmt='%10.8f')
        print("--> saved {}".format(out_fname))

    # 5) compare chi2
    # 需要注意将14位“转成”16位，与bin2fits保持一致
    pv_old = ( d[:,0:n_ref].sum(axis=1) - d[:,n_ref:2*n_ref].sum(axis=1) ) / n_ref
    pv_new = np.matmul(d,w[:-2])
    pv_old = (pv_old*4) #.astype(np.uint16)    
    pv_new = (pv_new*4) #.astype(np.uint16)
    std_old = np.std(pv_old)
    std_new = np.std(pv_new)

    return std_old, std_new

###################################################################################################
if __name__ == '__main__':
    mask = DCDS_Mask()

    mask37 = mask.gen_mask(37)
    n_ref=37
    n_sig=37
    std_old_37 = []
    std_new_37 = []
    fig = plt.figure(figsize=(16,10))
    for i in range(16):
        print('> processing ch: {}'.format(i+1))
        os_wave_dir = '0620/OS'+str(i+1)
        wfile = 'w_os_' + str(n_ref) + '_' + str(i+1)+'.txt'
        std_oldx, std_newx = getWeights(os_wave_dir, file_num=89, mask=mask37, n_ref=n_ref, n_sig=n_sig, out_fname=wfile)
        std_old_37.append(std_oldx)
        std_new_37.append(std_newx)
        tmp = np.loadtxt(wfile)
        plt.subplot(4,4,i+1)
        plt.plot(tmp,'--.')
        plt.title('ch :'+str(i+1))
    
    plt.tight_layout()
    plt.savefig('w_'+str(n_ref)+'.png')
    # np.savetxt('std_old_37.txt', np.array(std_old))
    # np.savetxt('std_new_37.txt', np.array(std_new))


    mask41 = mask.gen_mask(41)
    n_ref=41
    n_sig=41
    std_old_41 = []
    std_new_41 = []
    fig = plt.figure(figsize=(16,10))
    for i in range(16):
        print('> processing ch: {}'.format(i+1))
        os_wave_dir = '0620/OS'+str(i+1)
        wfile = 'w_os_' + str(n_ref) + '_' + str(i+1)+'.txt'
        std_oldx, std_newx = getWeights(os_wave_dir, file_num=89, mask=mask41, n_ref=n_ref, n_sig=n_sig, out_fname=wfile)
        std_old_41.append(std_oldx)
        std_new_41.append(std_newx)
        tmp = np.loadtxt(wfile)
        plt.subplot(4,4,i+1)
        plt.plot(tmp,'--.')
        plt.title('ch :'+str(i+1))
    
    plt.tight_layout()
    plt.savefig('w_'+str(n_ref)+'.png')
    # np.savetxt('std_old_41.txt', np.array(std_old))
    # np.savetxt('std_new_41.txt', np.array(std_new))
    
    mask45 = mask.gen_mask(45)
    n_ref=45
    n_sig=45
    std_old_45 = []
    std_new_45 = []
    fig = plt.figure(figsize=(16,10))
    for i in range(16):
        print('> processing ch: {}'.format(i+1))
        os_wave_dir = '0620/OS'+str(i+1)
        wfile = 'w_os_' + str(n_ref) + '_' + str(i+1)+'.txt'
        std_oldx, std_newx = getWeights(os_wave_dir, file_num=89, mask=mask45, n_ref=n_ref, n_sig=n_sig, out_fname=wfile)
        std_old_45.append(std_oldx)
        std_new_45.append(std_newx)
        tmp = np.loadtxt(wfile)
        plt.subplot(4,4,i+1)
        plt.plot(tmp,'--.')
        plt.title('ch :'+str(i+1))
    
    plt.tight_layout()
    plt.savefig('w_'+str(n_ref)+'.png')
    # np.savetxt('std_old_45.txt', np.array(std_old))
    # np.savetxt('std_new_45.txt', np.array(std_new))

    plt.figure()
    plt.plot(std_old_37, 'r:o', label='std_old_37')
    plt.plot(std_old_41, 'g:d', label='std_old_41')
    plt.plot(std_old_45, 'b:s', label='std_old_45')

    plt.plot(std_new_37, 'r--o', label='std_new_37')
    plt.plot(std_new_41, 'g--d', label='std_new_41')
    plt.plot(std_new_45, 'b--s', label='std_new_45')

    plt.legend()
    plt.show()
